{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "intro-docs-1",
   "metadata": {},
   "source": [
    "# Basic demonstration of Dspy integration functionality with the Unity Catalog AI Toolkit SDK\n",
    "\n",
    "To get started with this, you will need an OpenAI API Key. \n",
    "\n",
    "Once you have acquired your key, set it to the environment variable `OPENAI_API_KEY` after storing it in the `Databricks Secrets` API (accessible via dbutils or the databricks sdk workspace client)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fa44a8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -Uqqq unitycatalog-dspy[databricks] dspy mlflow\n",
    "\n",
    "%restart_python"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc9500c9",
   "metadata": {},
   "source": [
    "## Setting your API Key\n",
    "\n",
    "Don't forget to remove the key after you're done running cell 4!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d313bb8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64\n",
    "import os\n",
    "\n",
    "from databricks.sdk import WorkspaceClient\n",
    "\n",
    "workspace_client = WorkspaceClient()\n",
    "\n",
    "secret_scope = \"puneet_jain\"  # Change me!\n",
    "\n",
    "# Run this if you don't have the API key set to your secrets scope yet\n",
    "\n",
    "# if secret_scope not in [scope.name for scope in workspace_client.secrets.list_scopes()]:\n",
    "#     workspace_client.secrets.create_scope(secret_scope)\n",
    "\n",
    "# my_secret = \"<your API key, temporarily>\"\n",
    "\n",
    "# workspace_client.secrets.put_secret(scope=secret_scope, key=\"openai_api_key\", string_value=my_secret)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "api-key-note-1",
   "metadata": {},
   "source": [
    "## Fetch the key and set it to the environment variable key that the OpenAI SDK needs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "7b6d45e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"OPENAI_API_KEY\"] = base64.b64decode(\n",
    "    workspace_client.secrets.get_secret(scope=secret_scope, key=\"openai_api_key\").value\n",
    ").decode()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "8a674041",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert \"OPENAI_API_KEY\" in os.environ, (\n",
    "    \"Please set the OPENAI_API_KEY environment variable to your OpenAI API key\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "e2b14356",
   "metadata": {},
   "outputs": [],
   "source": [
    "from unitycatalog.ai.core.databricks import DatabricksFunctionClient\n",
    "from unitycatalog.ai.dspy.toolkit import UCFunctionToolkit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69590e6e",
   "metadata": {},
   "source": [
    "## Import the UC client for Databricks UC. \n",
    "This will allow for function creation through either the `create_function` API (requires the defined `sql_body` statement) or the `create_python_function` (requires a type-hint-applied and docstring commented python callable). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "97c52c4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "CATALOG = \"mlops_pj\"  # Change me !\n",
    "SCHEMA = \"rr_rag_chatbot\"  # Change me if you want"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "2fdf5b05",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = DatabricksFunctionClient()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a182f54a",
   "metadata": {},
   "source": [
    "### Define a Callable\n",
    "The requirements for the callable:\n",
    "\n",
    "**typing**\n",
    "\n",
    "Types **must** be supplied for both the arguments and the return type. Function signatures that do not have these defined will raise a `ValueError`.\n",
    "\n",
    "The following types are not allowed:\n",
    "`Union`\n",
    "`Any`\n",
    "\n",
    "Additional caveats:\n",
    "Collections **must** supply typing of the interior components. For instance, ``typing.Dict`` is not allowed, but ``typing.Dict[str, str]`` will work correctly. \n",
    "\n",
    "**doc strings**\n",
    "\n",
    "The doc string **must** be in the Google Docstring format.\n",
    "Args and Returns comments are optional, but the function description **is required**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "043fc753",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_weather(city: str) -> str:\n",
    "    \"\"\"Retrieve mock weather information for a given city.\n",
    "\n",
    "    This function looks up predefined mock weather data for a set of cities.\n",
    "    If the city is not found in the dataset, a default message is returned.\n",
    "\n",
    "    Args:\n",
    "        city (str): The name of the city to retrieve weather data for.\n",
    "\n",
    "    Returns:\n",
    "        str: A string describing the weather for the given city, or\n",
    "        \"Weather data not available\" if the city is not in the dataset.\n",
    "\n",
    "    Example:\n",
    "        >>> get_weather(\"New York\")\n",
    "        'Sunny, 25°C'\n",
    "\n",
    "        >>> get_weather(\"Boston\")\n",
    "        'Weather data not available'\n",
    "    \"\"\"\n",
    "    mock_data = {\n",
    "        \"New York\": \"Sunny, 25°C\",\n",
    "        \"Los Angeles\": \"Cloudy, 20°C\",\n",
    "        \"Chicago\": \"Rainy, 15°C\",\n",
    "        \"Houston\": \"Thunderstorms, 30°C\",\n",
    "        \"Phoenix\": \"Sunny, 35°C\",\n",
    "    }\n",
    "    return mock_data.get(city, \"Weather data not available\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e97351d",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.create_python_function(func=get_weather, catalog=CATALOG, schema=SCHEMA, replace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "c8237dc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a tool instance to use with Dspy\n",
    "\n",
    "toolkit = UCFunctionToolkit(function_names=[f\"{CATALOG}.{SCHEMA}.get_weather\"], client=client)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e55b2f4",
   "metadata": {},
   "source": [
    "## Enable tracing in MLflow\n",
    "Auto-enabling tracing allows us to see the calls made by dspy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc4c162f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlflow\n",
    "\n",
    "mlflow.dspy.autolog()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "5d2bcb4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dspy\n",
    "\n",
    "llm = dspy.LM(model=\"openai/gpt-4o\", api_key=os.getenv(\"OPENAI_API_KEY\"))\n",
    "dspy.configure(lm=llm)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d4ab4a3",
   "metadata": {},
   "source": [
    "### Build a simple ReAct agent\n",
    "We define a DSPy signature (`Agent_tools`) for the agent’s inputs/outputs and compose a ReAct module with `tools=toolkit.tools`. The agent maintains a minimal chat history and can call the UC function when helpful.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66a20a93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a ReAct agent\n",
    "react_agent = dspy.ReAct(signature=\"question -> answer\", tools=toolkit.tools, max_iters=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84d9c0dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AI:  The current weather in New York is Sunny with a temperature of 25°C. In Chicago, it is Rainy with a temperature of 15°C.\n"
     ]
    }
   ],
   "source": [
    "# Use the agent\n",
    "result = react_agent(question=\"What's the weather like in Tokyo?\")\n",
    "result.answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5cd8552",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get tool calls made\n",
    "result.trajectory"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
